import numpy as np
import os
import sys

if len(sys.argv)!=1:
    addJ=True
else:
    addJ=False

input=open('core','r').readline().split()
natom=int(input[0])
line=input[1].split('x')
paths=[int(line[0]),int(line[1]),int(line[2])]
addtype=bool(int(input[2]))
multi=paths[0]*paths[1]*paths[2]
N=natom*multi

def getdata(filename):
    if addJ:
        data=os.popen("grep 'atomic mag.' %s"%filename).readlines()
    else:
        data=os.popen("grep Tr %s"%filename).readlines()
    initial=data[0:N]
    step=data[N:2*N]
    final=data[2*N:]
    if len(initial) != len(final):
        open("error",'w').write("file error")
        sys.exit()
    if addtype:
        initial=[initial[0]]+initial[-multi+1:]
        step=[step[0]]+step[-multi+1:]
        final=[final[0]]+final[-multi+1:]
    else:
        initial=initial[:multi]
        step=step[:multi]
        final=final[:multi]
    disp_fs=[float(step[i].split()[-1]) for i in range(0,multi)]
    disp_fi=[float(final[i].split()[-1]) for i in range(0,multi)]
    return [disp_fs,disp_fi]

def fit_LS(data,Y,rank):
    if len(data) <= rank+1:
        print("Insufficient data")
        return
    X=[]
    tmp=[]
    for i in range(0,len(data)):
        for j in range(0,rank+1):
            tmp.append(data[i]**j)
        X.append(tmp)
        tmp=[]
    X=np.mat(X)
    Y=np.mat(Y).T
    C=(X.T*X).I    
    guess=C*X.T*Y
    error=(Y-X*guess).T.tolist()[0]
    stde=np.sqrt(sum([i**2 for i in error])/(len(data)-rank-1))
    trans=[C[i,i] for i in range(0,rank+1)]
    trans_error=[stde*np.sqrt(i) for i in trans]
    guess=guess.T.tolist()[0]    
    return [guess,trans_error]

def gallery():
    filelist=[i.rstrip() for i in os.popen("ls ./").readlines()]
    data={}
    Xdis=[]
    for outf in filelist:
        if outf[-4:] == ".out":
            dis=float(outf[outf.index('_')+1:outf.index('.out')])
            data[dis]=getdata(outf)
            Xdis.append(dis)
    Upass_fs=[]
    Upass_fi=[]
    for i in range(0,multi):
        pdata_fs=[data[j][0][i] for j in Xdis]
        pdata_fi=[data[j][1][i] for j in Xdis]
        Upass_fs.append(fit_LS(Xdis,pdata_fs,1)[0][1])
        Upass_fi.append(fit_LS(Xdis,pdata_fi,1)[0][1])
    return(Upass_fs,Upass_fi)

def relocate(initial,final):
    initial=[int(i) for i in initial.split('x')]
    final=[int(i) for i in final.split('x')]
    x=1+final[0]-initial[0]
    y=1+final[1]-initial[1]
    z=1+final[2]-initial[2]
    if x<=0:
        x+=paths[0]
    elif x>paths[0]:
        x-=paths[0]
    if y<=0:
        y+=paths[1]
    elif y>paths[1]:
        y-=paths[1]
    if z<=0:
        z+=paths[2]
    elif z>paths[2]:
        z-=paths[2]
    return "%sx%sx%s"%(x,y,z)

def reshape(Upass):
    data={}
    mat_Upass=[Upass]
    pos=[]
    for k in range(1,paths[2]+1):
        for j in range(1,paths[1]+1):
            for i in range(1,paths[0]+1):
                pos.append("%sx%sx%s"%(i,j,k))
    for i in range(0,multi):
        data[pos[i]]=Upass[i]
    Upline=[]
    for i in range(1,multi):
        for j in range(0,multi):
            Upline.append(data[relocate(pos[i],pos[j])])
        mat_Upass.append(Upline)
        Upline=[]
    return mat_Upass

Upass_data=gallery()
Upass_fstep=reshape(Upass_data[0])
Upass_fin=reshape(Upass_data[1])
Upass_disp=np.mat(Upass_fstep).I-np.mat(Upass_fin).I
if addJ:
    Upass_disp*=-2
print(Upass_disp)

#gallery()
#print(relocate("1x2","2x1"))
#reshape([1,2,3,4,5,6,7,8,9])
